package com.powerbank.gateway.filter;

import cn.hutool.json.JSONUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Mono;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 自定义限流过滤器
 * 支持多种限流策略：IP限流、用户限流、接口限流
 */
@Slf4j
@Component
public class CustomRateLimitGatewayFilterFactory extends AbstractGatewayFilterFactory<CustomRateLimitGatewayFilterFactory.Config> {

    @Autowired
    private StringRedisTemplate redisTemplate;

    // Lua脚本实现滑动窗口限流
    private static final String LIMIT_LUA_SCRIPT = 
        "local key = KEYS[1]\n" +
        "local window = tonumber(ARGV[1])\n" +
        "local limit = tonumber(ARGV[2])\n" +
        "local current = tonumber(redis.call('get', key) or 0)\n" +
        "if current < limit then\n" +
        "  redis.call('incr', key)\n" +
        "  redis.call('expire', key, window)\n" +
        "  return {current + 1, limit}\n" +
        "else\n" +
        "  return {current, limit}\n" +
        "end";

    private final DefaultRedisScript<List> limitScript;

    public CustomRateLimitGatewayFilterFactory() {
        super(Config.class);
        this.limitScript = new DefaultRedisScript<>();
        this.limitScript.setScriptText(LIMIT_LUA_SCRIPT);
        this.limitScript.setResultType(List.class);
    }

    @Override
    public GatewayFilter apply(Config config) {
        return (exchange, chain) -> {
            String path = exchange.getRequest().getPath().value();
            String method = exchange.getRequest().getMethod().name();
            
            // 获取限流键
            String rateLimitKey = getRateLimitKey(exchange, config);
            
            log.debug("限流检查: path={}, method={}, key={}, limit={}", 
                     path, method, rateLimitKey, config.getLimit());
            
            try {
                // 执行限流检查
                List<Long> results = redisTemplate.execute(
                    limitScript,
                    Arrays.asList(rateLimitKey),
                    String.valueOf(config.getWindowSize()),
                    String.valueOf(config.getLimit())
                );
                
                if (results != null && results.size() >= 2) {
                    Long current = results.get(0);
                    Long limit = results.get(1);
                    
                    // 添加限流响应头
                    exchange.getResponse().getHeaders().add("X-RateLimit-Limit", String.valueOf(limit));
                    exchange.getResponse().getHeaders().add("X-RateLimit-Remaining", String.valueOf(Math.max(0, limit - current)));
                    exchange.getResponse().getHeaders().add("X-RateLimit-Reset", String.valueOf(System.currentTimeMillis() + config.getWindowSize() * 1000));
                    
                    if (current > limit) {
                        log.warn("触发限流: key={}, current={}, limit={}", rateLimitKey, current, limit);
                        return rateLimited(exchange.getResponse(), "请求过于频繁，请稍后重试");
                    }
                }
                
                return chain.filter(exchange);
                
            } catch (Exception e) {
                log.error("限流检查异常: {}", e.getMessage(), e);
                // 限流检查失败时，允许请求通过，避免影响正常业务
                return chain.filter(exchange);
            }
        };
    }

    /**
     * 获取限流键
     */
    private String getRateLimitKey(org.springframework.web.server.ServerWebExchange exchange, Config config) {
        String keyPrefix = "rate_limit:" + config.getKeyType() + ":";
        
        switch (config.getKeyType()) {
            case "ip":
                String clientIp = getClientIp(exchange);
                return keyPrefix + clientIp;
            case "user":
                String userId = exchange.getRequest().getHeaders().getFirst("X-User-Id");
                return keyPrefix + (userId != null ? userId : "anonymous");
            case "path":
                String path = exchange.getRequest().getPath().value();
                return keyPrefix + path.replaceAll("/", "_");
            case "api":
                String api = exchange.getRequest().getPath().value() + ":" + exchange.getRequest().getMethod().name();
                return keyPrefix + api.replaceAll("/", "_");
            default:
                return keyPrefix + "default";
        }
    }

    /**
     * 获取客户端IP
     */
    private String getClientIp(org.springframework.web.server.ServerWebExchange exchange) {
        String clientIp = exchange.getRequest().getHeaders().getFirst("X-Forwarded-For");
        if (clientIp == null || clientIp.isEmpty() || "unknown".equalsIgnoreCase(clientIp)) {
            clientIp = exchange.getRequest().getHeaders().getFirst("X-Real-IP");
        }
        if (clientIp == null || clientIp.isEmpty() || "unknown".equalsIgnoreCase(clientIp)) {
            clientIp = exchange.getRequest().getRemoteAddress() != null ?
                      exchange.getRequest().getRemoteAddress().getAddress().getHostAddress() : "";
        }
        
        // 处理多个IP的情况，取第一个
        if (clientIp != null && clientIp.contains(",")) {
            clientIp = clientIp.split(",")[0].trim();
        }
        
        return clientIp;
    }

    /**
     * 返回限流响应
     */
    private Mono<Void> rateLimited(ServerHttpResponse response, String message) {
        response.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
        response.getHeaders().add("Content-Type", "application/json;charset=UTF-8");

        Map<String, Object> result = new HashMap<>();
        result.put("code", 429);
        result.put("message", message);
        result.put("data", null);
        result.put("timestamp", System.currentTimeMillis());

        String body = JSONUtil.toJsonStr(result);
        DataBuffer buffer = response.bufferFactory().wrap(body.getBytes(StandardCharsets.UTF_8));
        
        return response.writeWith(Mono.just(buffer));
    }

    /** 
     * 配置类
     */
    public static class Config {
        /** 
         * 限流类型：ip, user, path, api
         */
        private String keyType = "ip";
        
        /** 
         * 限流数量
         */
        private int limit = 100;
        
        /** 
         * 时间窗口大小（秒）
         */
        private int windowSize = 60;
        
        public String getKeyType() {
            return keyType;
        }
        
        public void setKeyType(String keyType) {
            this.keyType = keyType;
        }
        
        public int getLimit() {
            return limit;
        }
        
        public void setLimit(int limit) {
            this.limit = limit;
        }
        
        public int getWindowSize() {
            return windowSize;
        }
        
        public void setWindowSize(int windowSize) {
            this.windowSize = windowSize;
        }
    }
}